package indexer

import (
	"context"
	"fmt"
	"math"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"

	"ai-rag/config"

	"ai-rag/vectorstore"
	"github.com/cloudwego/ei
	"github.com/cloudwego/eino/components/indexer"
	"github.com/cloudwego/eino/schema"
	"ai-rag/vectorstore"
)

// Indexer 索引器接口
type Indexer interface {
	AddDocuments(ctx context.Context, docs []*schema.Document) error
	UpdateDocuments(ctx context.Context, docs []*schema.Document) error
	DeleteDocuments(ctx context.Context, ids []string) error
	Search(ctx context.Context, queryVector []float64, topK int) ([]*schema.Document, error)
	GetDocuments(ctx context.Context, ids []string) ([]*schema.Document, error)
	Save(ctx context.Context) error
	Load(ctx context.Context) error
	Close() error
}

// GenericIndexer 通用索引器实现
	typeName    string
	indexPath   string
	storagePath string
	useGPU      bool
	einoIndexer indexer.Indexer
	einoIndexer  indexer.Indexer
}

// NewGenericIndexer 创建新的通用索引器
func NewGenericIndexer(cfg *config.IndexConfig) (*GenericIndexer, error) {
	// 确保目录存在
	if err := os.MkdirAll(cfg.IndexPath, 0755); err != nil {
		return nil, fmt.Errorf("创建索引目录失败: %w", err)
	}

	if err := os.MkdirAll(cfg.StoragePath, 0755); err != nil {
		return nil, fmt.Errorf("创建存储目录失败: %w", err)
	}

	var einoIndexer Indexer
	var err error

	switch cfg.Type {
	case "faiss":
		// 使用vectorstore包中的FAISS索引器
		vectorstoreConfig := &vectorstore.FaissIndexerConfig{
			IndexPath:  filepath.Join(cfg.IndexPath, "faiss.index"),
			Dimension:  cfg.Dimension,
			MetricType: cfg.MetricType,
		}
		vectorIndexer, err := vectorstore.NewFaissIndexer(vectorstoreConfig)
		if err != nil {
			return nil, err
		}
		einoIndexer = &VectorStoreAdapter{indexer: vectorIndexer}
	case "chroma":
		// 使用vectorstore包中的Chroma索引器
			DBPath:         cfg.IndexPath,
			DBPath:        cfg.IndexPath,
			Dimension:      cfg.Dimension,
			MetricType:     cfg.MetricType,
			MetricType:    cfg.MetricType,
		}
		vectorIndexer, err := vectorstore.NewChromaIndexer(vectorstoreConfig)
		if err != nil {
			return nil, err
		}
		einoIndexer = &VectorStoreAdapter{indexer: vectorIndexer}
	case "milvus":
		// 使用vectorstore包中的Milvus索引器
		// 解析地址为host和port
		host := cfg.Address
		port := "19530" // 默认端口
		if strings.Contains(cfg.Address, ":") {
			parts := strings.Split(cfg.Address, ":")
			host = parts[0]
			if len(parts) > 1 {
				port = parts[1]
			}
		}
			Host:           host,
			Port:           port,
			Port:          port,
			Dimension:      cfg.Dimension,
			MetricType:     cfg.MetricType,
			MetricType:    cfg.MetricType,
		}
		vectorIndexer, err := vectorstore.NewMilvusIndexer(vectorstoreConfig)
		if err != nil {
			return nil, err
		}
		einoIndexer = &VectorStoreAdapter{indexer: vectorIndexer}
	case "pinecone":
		// 使用vectorstore包中的Pinecone索引器
		apiKey := getEnvOrDefault("PINECONE_API_KEY", cfg.APIKey)
		vectorstoreConfig := &vectorstore.PineconeIndexerConfig{
			ApiKey:      apiKey,
			Environment: cfg.Environment,
			IndexName:   cfg.IndexName,
			Namespace:   "", // 可选参数，根据需求设置
			Dimension:   cfg.Dimension,
			MetricType:  cfg.MetricType,
		}
		vectorIndexer, err := vectorstore.NewPineconeIndexer(vectorstoreConfig)
		if err != nil {
			return nil, err
		}
		einoIndexer = &VectorStoreAdapter{indexer: vectorIndexer}
	case "memory":
		// 内存索引器，用于开发和测试
		einoIndexer = &MemoryIndexer{
			documents: make(map[string]*schema.Document),
			dimension: cfg.Dimension,
		}
	default:
		return nil, fmt.Errorf("未实现的索引类型: %s", cfg.Type)
	}

	if err != nil {
		return nil, fmt.Errorf("初始化%s索引器失败: %w", cfg.Type, err)
	}

		typeName:    cfg.Type,
		indexPath:   cfg.IndexPath,
		storagePath: cfg.StoragePath,
		useGPU:      cfg.UseGPU,
		einoIndexer: einoIndexer,
		einoIndexer:  einoIndexer,
	}, nil
}

// getEnvOrDefault 获取环境变量，如果不存在则返回默认值
func getEnvOrDefault(key, defaultValue string) string {
	if value := os.Getenv(key); value != "" {
		return value
	}
	return defaultValue
}

// AddDocuments 添加文档到索引
func (idx *GenericIndexer) AddDocuments(ctx context.Context, docs []*schema.Document) error {
	logrus.WithField("count", len(docs)).Info("开始添加文档到索引")

	// 验证所有文档是否都有嵌入向量
	for i, doc := range docs {
		if doc.Embedding == nil || len(doc.Embedding) == 0 {
			return fmt.Errorf("文档 %d 没有嵌入向量", i)
		}
		// 为没有ID的文档生成唯一ID
		if doc.ID == "" {
			doc.ID = generateDocumentID(doc)
		}
	}

	if err := idx.einoIndexer.AddDocuments(ctx, docs); err != nil {
		logrus.WithError(err).Error("添加文档到索引失败")
		return err
	}

	logrus.WithField("count", len(docs)).Info("文档添加到索引成功")
	return nil
}

// UpdateDocuments 更新索引中的文档
func (idx *GenericIndexer) UpdateDocuments(ctx context.Context, docs []*schema.Document) error {
	logrus.WithField("count", len(docs)).Info("开始更新索引中的文档")

	// 验证所有文档是否都有ID和嵌入向量
	for i, doc := range docs {
		if doc.ID == "" {
			return fmt.Errorf("文档 %d 没有ID，无法更新", i)
		}
		if doc.Embedding == nil || len(doc.Embedding) == 0 {
			return fmt.Errorf("文档 %d 没有嵌入向量", i)
		}
	}

	// 尝试先删除再添加的方式更新文档
	ids := make([]string, len(docs))
	for i, doc := range docs {
		ids[i] = doc.ID
	}

	if updater, ok := idx.einoIndexer.(interface {
		UpdateDocuments(context.Context, []*schema.Document) error
	}); ok {
	if updater, ok := idx.einoIndexer.(interface{ UpdateDocuments(context.Context, []*schema.Document) error }); ok {
		if err := updater.UpdateDocuments(ctx, docs); err != nil {
			logrus.WithError(err).Error("更新文档失败")
			return err
		}
	} else {
		// 否则使用删除再添加的方式
		if deleter, ok := idx.einoIndexer.(interface {
			DeleteDocuments(context.Context, []string) error
		}); ok {
		if deleter, ok := idx.einoIndexer.(interface{ DeleteDocuments(context.Context, []string) error }); ok {
			_ = deleter.DeleteDocuments(ctx, ids)
		}
		// 再添加新文档
		if err := idx.einoIndexer.AddDocuments(ctx, docs); err != nil {
			logrus.WithError(err).Error("添加更新后的文档失败")
			return err
		}
	}

	logrus.WithField("count", len(docs)).Info("文档更新成功")
	return nil
}

// DeleteDocuments 从索引中删除文档
func (idx *GenericIndexer) DeleteDocuments(ctx context.Context, ids []string) error {
	logrus.WithField("count", len(ids)).Info("开始从索引中删除文档")

	if len(ids) == 0 {
		logrus.Warn("删除文档列表为空")
		return nil
	}

	if deleter, ok := idx.einoIndexer.(interface {
		DeleteDocuments(context.Context, []string) error
	}); ok {
	if deleter, ok := idx.einoIndexer.(interface{ DeleteDocuments(context.Context, []string) error }); ok {
		if err := deleter.DeleteDocuments(ctx, ids); err != nil {
			logrus.WithError(err).Error("删除文档失败")
			return err
		}
	} else {
		// 如果不支持，记录警告但不返回错误
		logrus.Warn("当前索引器不支持删除文档功能")
	}

	logrus.WithField("count", len(ids)).Info("文档删除成功")
	return nil
}

// GetDocuments 根据ID获取文档
func (idx *GenericIndexer) GetDocuments(ctx context.Context, ids []string) ([]*schema.Document, error) {
	logrus.WithField("count", len(ids)).Info("开始根据ID获取文档")

	if len(ids) == 0 {
		return []*schema.Document{}, nil
	}

	if getter, ok := idx.einoIndexer.(interface {
		GetDocuments(context.Context, []string) ([]*schema.Document, error)
	}); ok {
	if getter, ok := idx.einoIndexer.(interface{ GetDocuments(context.Context, []string) ([]*schema.Document, error) }); ok {
		docs, err := getter.GetDocuments(ctx, ids)
		if err != nil {
			logrus.WithError(err).Error("获取文档失败")
			return nil, err
		}
		logrus.WithField("found", len(docs)).Info("获取文档成功")
		return docs, nil
	}

	// 如果不支持，返回一个空列表
	logrus.Warn("当前索引器不支持根据ID获取文档功能")
	return []*schema.Document{}, nil
}

// generateDocumentID 为文档生成唯一ID
func generateDocumentID(doc *schema.Document) string {
	// 使用文档内容的前100个字符和时间戳生成ID
	content := doc.Content
	if len(content) > 100 {
		content = content[:100]
	}
	// 移除特殊字符
	content = strings.ReplaceAll(content, "\n", " ")
	content = strings.ReplaceAll(content, "\r", " ")
	// 使用内容的哈希值和时间戳生成ID
	return fmt.Sprintf("doc_%x_%d", []byte(content)[:10], time.Now().UnixNano())
}

// Search 在索引中搜索相似文档
func (idx *GenericIndexer) Search(ctx context.Context, queryVector []float64, topK int) ([]*schema.Document, error) {
	logrus.WithField("topK", topK).Info("开始在索引中搜索相似文档")

	if queryVector == nil || len(queryVector) == 0 {
		return nil, fmt.Errorf("查询向量不能为空")
	}

	docs, err := idx.einoIndexer.Search(ctx, queryVector, topK)
	if err != nil {
		logrus.WithError(err).Error("搜索文档失败")
		return nil, err
	}

	logrus.WithField("found", len(docs)).Info("文档搜索完成")
	return docs, nil
}

// Save 保存索引
func (idx *GenericIndexer) Save(ctx context.Context) error {
	logrus.Info("开始保存索引")

	if err := idx.einoIndexer.Save(ctx); err != nil {
		logrus.WithError(err).Error("保存索引失败")
		return err
	}

	logrus.Info("索引保存成功")
	return nil
}

// Load 加载索引
func (idx *GenericIndexer) Load(ctx context.Context) error {
	logrus.Info("开始加载索引")

	// 检查底层索引器是否支持Load方法
	if loader, ok := idx.einoIndexer.(interface{ Load(context.Context) error }); ok {
		if err := loader.Load(ctx); err != nil {
			// 如果是文件不存在错误，可以返回nil，表示索引为空
			if os.IsNotExist(err) {
				logrus.Info("索引文件不存在，将创建新索引")
				return nil
			}
			logrus.WithError(err).Error("加载索引失败")
			return err
		}
	}

	logrus.Info("索引加载成功")
	return nil
}

// Close 关闭索引器
func (idx *GenericIndexer) Close() error {
	logrus.Info("开始关闭索引器")

	// 检查底层索引器是否支持Close方法
	if closer, ok := idx.einoIndexer.(interface{ Close() error }); ok {
		if err := closer.Close(); err != nil {
			logrus.WithError(err).Error("关闭索引器失败")
			return err
		}
	}

	logrus.Info("索引器关闭成功")
	return nil
}

// MemoryIndexer 内存索引器实现，用于开发和测试
// 这是一个简单的内存索引器，不依赖外部服务

type MemoryIndexer struct {
	documents map[string]*schema.Document
	dimension int
}

// AddDocuments 添加文档到内存索引
func (m *MemoryIndexer) AddDocuments(ctx context.Context, docs []*schema.Document) error {
	for _, doc := range docs {
		m.documents[doc.ID] = doc
	}
	return nil
}

// Search 在内存索引中搜索相似文档
func (m *MemoryIndexer) Search(ctx context.Context, queryVector []float64, topK int) ([]*schema.Document, error) {
	// 简单的余弦相似度搜索实现
	type scoredDoc struct {
		doc   *schema.Document
		score float64
	}

	var scoredDocs []scoredDoc
	for _, doc := range m.documents {
		if doc.Embedding == nil || len(doc.Embedding) == 0 {
			continue
		}
		// 计算余弦相似度
		score := cosineSimilarity(queryVector, doc.Embedding)
		scoredDocs = append(scoredDocs, scoredDoc{doc: doc, score: score})
	}

	// 按相似度排序
	sort.Slice(scoredDocs, func(i, j int) bool {
		return scoredDocs[i].score > scoredDocs[j].score
	})

	// 返回前topK个结果
	resultCount := topK
	if resultCount > len(scoredDocs) {
		resultCount = len(scoredDocs)
	}

	results := make([]*schema.Document, resultCount)
	for i := 0; i < resultCount; i++ {
		results[i] = scoredDocs[i].doc
	}

	return results, nil
}

// Save 保存索引（内存索引不需要保存）
func (m *MemoryIndexer) Save(ctx context.Context) error {
	return nil
}

// Load 加载索引（内存索引不需要加载）
func (m *MemoryIndexer) Load(ctx context.Context) error {
	return nil
}

// Close 关闭索引器
func (m *MemoryIndexer) Close() error {
	return nil
}

// DeleteDocuments 从内存索引中删除文档
func (m *MemoryIndexer) DeleteDocuments(ctx context.Context, ids []string) error {
	for _, id := range ids {
		delete(m.documents, id)
	}
	return nil
}

// GetDocuments 从内存索引中获取文档
func (m *MemoryIndexer) GetDocuments(ctx context.Context, ids []string) ([]*schema.Document, error) {
	var results []*schema.Document
	for _, id := range ids {
		if doc, exists := m.documents[id]; exists {
			results = append(results, doc)
		}
	}
	return results, nil
}

// VectorStoreAdapter 适配vectorstore包中的索引器到indexer包的接口
type VectorStoreAdapter struct {
	indexer vectorstore.Indexer
}

// AddDocuments 添加文档到索引
func (v *VectorStoreAdapter) AddDocuments(ctx context.Context, docs []*schema.Document) error {
	return v.indexer.AddDocuments(ctx, docs)
}

// UpdateDocuments 更新文档
func (v *VectorStoreAdapter) UpdateDocuments(ctx context.Context, docs []*schema.Document) error {
	return v.indexer.UpdateDocuments(ctx, docs)
}

// DeleteDocuments 删除文档
func (v *VectorStoreAdapter) DeleteDocuments(ctx context.Context, ids []string) error {
	return v.indexer.DeleteDocuments(ctx, ids)
}

// Search 搜索相似文档
func (v *VectorStoreAdapter) Search(ctx context.Context, queryVector []float64, topK int) ([]*schema.Document, error) {
	return v.indexer.Search(ctx, queryVector, topK)
}

// GetDocuments 根据ID获取文档
func (v *VectorStoreAdapter) GetDocuments(ctx context.Context, ids []string) ([]*schema.Document, error) {
	return v.indexer.GetDocuments(ctx, ids)
}

// Save 保存索引
func (v *VectorStoreAdapter) Save(ctx context.Context) error {
	return v.indexer.Save(ctx)
}

// Load 加载索引
func (v *VectorStoreAdapter) Load(ctx context.Context) error {
	return v.indexer.Load(ctx)
}

// Close 关闭索引器
func (v *VectorStoreAdapter) Close() error {
	return v.indexer.Close()
}

// cosineSimilarity 计算两个向量的余弦相似度
func cosineSimilarity(a, b []float64) float64 {
	if len(a) != len(b) {
		return 0
	}

	dotProduct := 0.0
	normA := 0.0
	normB := 0.0

	for i := range a {
		dotProduct += a[i] * b[i]
		normA += a[i] * a[i]
		normB += b[i] * b[i]
	}

	if normA == 0 || normB == 0 {
		return 0
	}

	return dotProduct / (math.Sqrt(normA) * math.Sqrt(normB))
}

// BatchIndexer 批量索引器，用于高效处理大量文档
	indexer   Indexer
	batchSize int
	batchSize  int
}

// NewBatchIndexer 创建新的批量索引器
func NewBatchIndexer(indexer Indexer, batchSize int) *BatchIndexer {
	return &BatchIndexer{
		indexer:   indexer,
		batchSize: batchSize,
	}
}

// AddDocuments 批量添加文档到索引
func (bi *BatchIndexer) AddDocuments(ctx context.Context, docs []*schema.Document) error {
	totalCount := len(docs)
	logrus.WithField("total_count", totalCount).WithField("batch_size", bi.batchSize).Info("开始批量添加文档")

	for i := 0; i < totalCount; i += bi.batchSize {
		end := i + bi.batchSize
		if end > totalCount {
			end = totalCount
		}

		batch := docs[i:end]
		logrus.WithField("batch_start", i).WithField("batch_end", end-1).Info("处理文档批次")

		if err := bi.indexer.AddDocuments(ctx, batch); err != nil {
			return fmt.Errorf("添加批次文档失败 (批次 %d-%d): %w", i, end-1, err)
		}

		// 每处理一个批次就保存一次索引
		if err := bi.indexer.Save(ctx); err != nil {
			logrus.WithError(err).Warn("保存索引失败，但继续处理")
		}

		logrus.WithField("processed", end).WithField("total", totalCount).Info("文档索引进度")
	}

	// 最后再保存一次索引，确保所有文档都被保存
	if err := bi.indexer.Save(ctx); err != nil {
		logrus.WithError(err).Error("最终保存索引失败")
		return err
	}

	logrus.WithField("total_count", totalCount).Info("批量文档索引完成")
	return nil
}

// UpdateDocuments 批量更新文档
func (bi *BatchIndexer) UpdateDocuments(ctx context.Context, docs []*schema.Document) error {
	totalCount := len(docs)
	logrus.WithField("total_count", totalCount).WithField("batch_size", bi.batchSize).Info("开始批量更新文档")

	for i := 0; i < totalCount; i += bi.batchSize {
		end := i + bi.batchSize
		if end > totalCount {
			end = totalCount
		}

		batch := docs[i:end]
		logrus.WithField("batch_start", i).WithField("batch_end", end-1).Info("处理文档批次")

		if err := bi.indexer.UpdateDocuments(ctx, batch); err != nil {
			return fmt.Errorf("更新批次文档失败 (批次 %d-%d): %w", i, end-1, err)
		}

		// 每处理一个批次就保存一次索引
		if err := bi.indexer.Save(ctx); err != nil {
			logrus.WithError(err).Warn("保存索引失败，但继续处理")
		}

		logrus.WithField("processed", end).WithField("total", totalCount).Info("文档更新进度")
	}

	// 最后再保存一次索引，确保所有文档都被保存
	if err := bi.indexer.Save(ctx); err != nil {
		logrus.WithError(err).Error("最终保存索引失败")
		return err
	}

	logrus.WithField("total_count", totalCount).Info("批量文档更新完成")
	return nil
}

// DeleteDocuments 批量删除文档
func (bi *BatchIndexer) DeleteDocuments(ctx context.Context, ids []string) error {
	totalCount := len(ids)
	logrus.WithField("total_count", totalCount).WithField("batch_size", bi.batchSize).Info("开始批量删除文档")

	for i := 0; i < totalCount; i += bi.batchSize {
		end := i + bi.batchSize
		if end > totalCount {
			end = totalCount
		}

		batch := ids[i:end]
		logrus.WithField("batch_start", i).WithField("batch_end", end-1).Info("处理文档ID批次")

		if err := bi.indexer.DeleteDocuments(ctx, batch); err != nil {
			return fmt.Errorf("删除批次文档失败 (批次 %d-%d): %w", i, end-1, err)
		}

		logrus.WithField("processed", end).WithField("total", totalCount).Info("文档删除进度")
	}

	logrus.WithField("total_count", totalCount).Info("批量文档删除完成")
	return nil
}

// GetDocuments 批量获取文档
func (bi *BatchIndexer) GetDocuments(ctx context.Context, ids []string) ([]*schema.Document, error) {
	totalCount := len(ids)
	logrus.WithField("total_count", totalCount).WithField("batch_size", bi.batchSize).Info("开始批量获取文档")

	var allDocs []*schema.Document

	for i := 0; i < totalCount; i += bi.batchSize {
		end := i + bi.batchSize
		if end > totalCount {
			end = totalCount
		}

		batch := ids[i:end]
		logrus.WithField("batch_start", i).WithField("batch_end", end-1).Info("处理文档ID批次")

		docs, err := bi.indexer.GetDocuments(ctx, batch)
		if err != nil {
			return nil, fmt.Errorf("获取批次文档失败 (批次 %d-%d): %w", i, end-1, err)
		}

		allDocs = append(allDocs, docs...)
		logrus.WithField("processed", end).WithField("total", totalCount).Info("文档获取进度")
	}

	logrus.WithField("total_found", len(allDocs)).Info("批量文档获取完成")
	return allDocs, nil
}

// Search 搜索文档
func (bi *BatchIndexer) Search(ctx context.Context, queryVector []float64, topK int) ([]*schema.Document, error) {
	return bi.indexer.Search(ctx, queryVector, topK)
}

// Save 保存索引
func (bi *BatchIndexer) Save(ctx context.Context) error {
	return bi.indexer.Save(ctx)
}

// Load 加载索引
func (bi *BatchIndexer) Load(ctx context.Context) error {
	return bi.indexer.Load(ctx)
}

// Close 关闭索引器
func (bi *BatchIndexer) Close() error {
	return bi.indexer.Close()
}